package com.gaga.console.core.service.impl.ai;

import com.gaga.common.enums.PolishOrQaType;
import com.gaga.common.result.AiModelApiResult;
import com.gaga.console.core.entity.dto.ChatDto;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

@Slf4j
public class DoubaoAiHandler extends AbstractAiHandler {

    // 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    // 此为默认路径，您可根据业务所在地域进行配置
    public static String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    public static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    public static Dispatcher dispatcher = new Dispatcher();

    @Override
    public AiModelApiResult getAiOutMessage(ChatDto dto) {
        AiModelApiResult aiModelApiResult = new AiModelApiResult();
        aiModelApiResult.setIsSuceess(false);

        String apiKey = dto.getToken();

        ArkService arkService = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();

        String promptMessage = buildPromptMessage(dto);

        List<ChatMessage> messages = new ArrayList<>();
        if (PolishOrQaType.QA.equals(dto.getPolishOrQaType())) {
            String liveRoomDesc = dto.getLiveRoomDesc();
            if (StringUtils.isNotBlank(dto.getAnchorName())) {
                liveRoomDesc = liveRoomDesc + "主播名字是" + dto.getAnchorName();
            }
            if (StringUtils.isNotBlank(dto.getLiveRoomDesc())) {
                liveRoomDesc = liveRoomDesc + "。用户是" + dto.getUserCallName();
            }
            ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(liveRoomDesc).build();
            messages.add(systemMessage);
            ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(promptMessage).build();
            messages.add(userMessage);
        } else {
            ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(polishPromptText + promptMessage).build();
            messages.add(userMessage);
        }

        Long startTime = System.currentTimeMillis();

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                // 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                .model("doubao-1-5-pro-32k-250115")
                .messages(messages)
                .temperature(0.7)
                .frequencyPenalty(0.0)
                .topP(0.7)
                .build();

//        log.info("doubao润色接口参数=" + JSON.toJSONString(messages));

        ChatCompletionResult chatCompletion = null;
        try {
            chatCompletion = arkService.createChatCompletion(chatCompletionRequest);
        } catch (Exception e) {
            log.error("doubao润色失败==={}", e.getMessage());
            aiModelApiResult.setResult(e.getMessage());
        }

        if (chatCompletion != null) {
            if (CollectionUtils.isNotEmpty(chatCompletion.getChoices())) {
                aiModelApiResult.setIsSuceess(true);
                aiModelApiResult.setResult(chatCompletion.getChoices().get(0).getMessage().getContent().toString());
                log.info("润色之后的话:" + aiModelApiResult.getResult());
            }
        }
        log.info("润色请求耗时: {}ms", System.currentTimeMillis() - startTime);

        arkService.shutdownExecutor();
        return aiModelApiResult;
    }
}
